import segmentation_models as sm
import cv2
import numpy as np
import matplotlib.pyplot as plt
import albumentations as A

# helper function for data visualization
def visualize(**images):
    """PLot images in one row."""
    n = len(images)
    plt.figure(figsize=(16, 5))
    for i, (name, image) in enumerate(images.items()):
        plt.subplot(1, n, i + 1)
        plt.xticks([])
        plt.yticks([])
        plt.title(' '.join(name.split('_')).title())
        plt.imshow(image)
    plt.show()

# helper function for data visualization
def denormalize(x):
    """Scale image to range 0..1 for correct plot"""
    x_max = np.percentile(x, 98)
    x_min = np.percentile(x, 2)
    x = (x - x_min) / (x_max - x_min)
    x = x.clip(0, 1)
    return x

def get_preprocessing(preprocessing_fn):
    """Construct preprocessing transform

    Args:
        preprocessing_fn (callbale): data normalization function
            (can be specific for each pretrained neural network)
    Return:
        transform: albumentations.Compose

    """

    _transform = [
        A.Lambda(image=preprocessing_fn),
    ]
    return A.Compose(_transform)

BACKBONE = 'efficientnetb3'
BATCH_SIZE = 8
CLASSES = ['car']
LR = 0.0001
EPOCHS = 10


testpath = '../data/CamVid/test/'
testfilename = '0001TP_008640.png'


preprocess_input = sm.get_preprocessing(BACKBONE)
preprocessf = get_preprocessing(preprocess_input)
image = cv2.imread(testpath+testfilename)
# image = preprocessf(force_apply=False, image=image)
# image = np.expand_dims(image, axis=0)


# define network parameters
n_classes = 1 if len(CLASSES) == 1 else (len(CLASSES) + 1)  # case for binary and multiclass segmentation
activation = 'sigmoid' if n_classes == 1 else 'softmax'


model = sm.Unet(BACKBONE, classes=n_classes, activation=activation)


model.load_weights('best_model.h5')



pr_mask = model.predict(image).round()

visualize(
    image=denormalize(image.squeeze()),
    pr_mask=pr_mask[..., 0].squeeze()
)

